{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SUknhuHKyc-E"
   },
   "source": [
    "# <center>OpenAI agent pattern: orchestrator and workers</center>\n",
    "\n",
    "A starter guide for building an agent loop using the `openai-agents` library.\n",
    "\n",
    "This pattern uses orchestators and workers. The orchestrator chooses which worker to use for a specific sub-task. The worker attempts to complete the sub-task and return a result. The orchestrator then uses the result to choose the next worker to use until a final result is returned.\n",
    "\n",
    "In the following example, we'll build an agent which creates a portfolio of stocks and ETFs based on a user's investment strategy.\n",
    "1.  **Orchestrator:** Chooses which worker to use based on the user's investment strategy.\n",
    "2.  **Research Agent:** Searches the web for information about stocks and ETFs that could support the user's investment strategy.\n",
    "3.  **Evaluation Agent:** Evaluates the research report and provides feedback on what data is missing.\n",
    "4.  **Portfolio Agent:** Creates a portfolio of stocks and ETFs based on the research report."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "n69HR7eJswNt"
   },
   "source": [
    "### Install Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install base libraries for OpenAI\n",
    "!pip install -q openai openai-agents\n",
    "\n",
    "# Install optional libraries for OpenInference/OpenTelemetry tracing\n",
    "!pip install -q arize-phoenix-otel openinference-instrumentation-openai-agents openinference-instrumentation-openai"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jQnyEnJisyn3"
   },
   "source": [
    "### Setup Keys\n",
    "\n",
    "Add your OpenAI API key to the environment variable `OPENAI_API_KEY`.\n",
    "\n",
    "Copy your Phoenix `API_KEY` from your settings page at [app.phoenix.arize.com](https://app.phoenix.arize.com)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "if \"OPENAI_API_KEY\" not in os.environ:\n",
    "    os.environ[\"OPENAI_API_KEY\"] = getpass(\"🔑 Enter your OpenAI API key: \")\n",
    "\n",
    "if \"PHOENIX_API_KEY\" not in os.environ:\n",
    "    os.environ[\"PHOENIX_API_KEY\"] = getpass(\"🔑 Enter your Phoenix API key: \")\n",
    "\n",
    "if \"PHOENIX_COLLECTOR_ENDPOINT\" not in os.environ:\n",
    "    os.environ[\"PHOENIX_COLLECTOR_ENDPOINT\"] = getpass(\"🔑 Enter your Phoenix Collector Endpoint\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kfid5cE99yN5"
   },
   "source": [
    "### Setup Tracing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from phoenix.otel import register\n",
    "\n",
    "tracer_provider = register(\n",
    "    project_name=\"openai-agents\",\n",
    "    auto_instrument=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating the agents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "from textwrap import dedent\n",
    "from typing import Literal\n",
    "\n",
    "from agents import Agent, Runner, TResponseInputItem, WebSearchTool\n",
    "from agents.model_settings import ModelSettings\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "\n",
    "class PortfolioItem(BaseModel):\n",
    "    ticker: str = Field(description=\"The ticker of the stock or ETF.\")\n",
    "    allocation: float = Field(\n",
    "        description=\"The percentage allocation of the ticker in the portfolio. The sum of all allocations should be 100.\"\n",
    "    )\n",
    "    reason: str = Field(description=\"The reason why this ticker is included in the portfolio.\")\n",
    "\n",
    "\n",
    "class Portfolio(BaseModel):\n",
    "    tickers: list[PortfolioItem] = Field(\n",
    "        description=\"A list of tickers that could support the user's stated investment strategy.\"\n",
    "    )\n",
    "\n",
    "\n",
    "class EvaluationFeedback(BaseModel):\n",
    "    feedback: str = Field(\n",
    "        description=\"What data is missing in order to create a portfolio of stocks and ETFs based on the user's investment strategy.\"\n",
    "    )\n",
    "    score: Literal[\"pass\", \"needs_improvement\", \"fail\"] = Field(\n",
    "        description=\"A score on the research report. Pass if you have at least 5 tickers with data that supports the user's investment strategy to create a portfolio, needs_improvement if you do not have enough supporting data, and fail if you have no tickers.\"\n",
    "    )\n",
    "\n",
    "\n",
    "evaluation_agent = Agent(\n",
    "    name=\"Evaluation Agent\",\n",
    "    instructions=dedent(\n",
    "        \"\"\"You are a senior financial analyst. You will be provided with a stock research report with positive and negative catalysts. Your task is to evaluate the report and provide feedback on what to improve.\"\"\"\n",
    "    ),\n",
    "    model=\"gpt-4.1\",\n",
    "    output_type=EvaluationFeedback,\n",
    ")\n",
    "\n",
    "portfolio_agent = Agent(\n",
    "    name=\"Portfolio Agent\",\n",
    "    instructions=dedent(\n",
    "        \"\"\"You are a senior financial analyst. You will be provided with a stock research report. Your task is to create a portfolio of stocks and ETFs that could support the user's stated investment strategy. Include facts and data from the research report in the stated reasons for the portfolio allocation.\"\"\"\n",
    "    ),\n",
    "    model=\"o4-mini\",\n",
    "    output_type=Portfolio,\n",
    ")\n",
    "\n",
    "research_agent = Agent(\n",
    "    name=\"FinancialSearchAgent\",\n",
    "    instructions=dedent(\n",
    "        \"\"\"You are a research assistant specializing in financial topics. Given a stock ticker, use web search to retrieve up‑to‑date context and produce a short summary of at most 50 words. Focus on key numbers, events, or quotes that will be useful to a financial analyst.\"\"\"\n",
    "    ),\n",
    "    model=\"gpt-4.1\",\n",
    "    tools=[WebSearchTool()],\n",
    "    model_settings=ModelSettings(tool_choice=\"required\", parallel_tool_calls=True),\n",
    ")\n",
    "\n",
    "orchestrator_agent = Agent(\n",
    "    name=\"Routing Agent\",\n",
    "    instructions=dedent(\"\"\"You are a senior financial analyst. You are trying to create a portfolio based on my stated investment strategy. Your task is to handoff to the appropriate agent or tool.\n",
    "\n",
    "    First, handoff to the research_agent to give you a report on stocks and ETFs that could support the user's stated investment strategy.\n",
    "    Then, handoff to the evaluation_agent to give you a score on the research report. If the evaluation_agent returns a needs_improvement or fail, continue using the research_agent to gather more information.\n",
    "    Once the evaluation_agent returns a pass, handoff to the portfolio_agent to create a portfolio.\"\"\"),\n",
    "    model=\"gpt-4.1\",\n",
    "    handoffs=[\n",
    "        research_agent,\n",
    "        evaluation_agent,\n",
    "        portfolio_agent,\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_input = input(\"Enter your investment strategy: \")\n",
    "input_items: list[TResponseInputItem] = [{\"content\": user_input, \"role\": \"user\"}]\n",
    "\n",
    "while True:\n",
    "    orchestrator = await Runner.run(orchestrator_agent, input_items)\n",
    "    orchestrator_output = orchestrator.final_output\n",
    "    pprint(orchestrator_output)\n",
    "\n",
    "    input_items = orchestrator.to_input_list()\n",
    "    if isinstance(orchestrator_output, Portfolio):\n",
    "        break\n",
    "    print(\"Going back to orchestrator\")\n",
    "    # input_items.append({\"content\": f\"Keep going\", \"role\": \"user\"})\n",
    "\n",
    "print(\"AGENT COMPLETE\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
